import cv2
import numpy as np
from siamrpn import TrackerSiamRPN
from criteria import iou

class tracker_manager:

    def __init__(self, params, track_method='SiamRPN'):

        """
        @ param tracker_state : =1  means the tracker is working well
                              : =-1 means the tracker should be removed
        @ param tracker_list : tracker elements created from CV, each links to the element of tracker_state list
        @ param tracker_bbox : tracker bbox's list
        @ param distance_threshold : if the distance between the key bbox and tracker bbox is bigger than the threshold,
                                     we regard the two bbox are not matched with each other
        @ param colors : the random color of each tracker bbox
        @ param key_bbox_size : the width and height of each key bbox, and using them to modify the tracker bbox's width and height
        """

        self.tracker_state = np.array([])
        self.tracker_list = []
        self.distance_threshold = 10
        self.overlap_threshold = 0.7
        color0 = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
        color1 = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
        color2 = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
        color3 = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
        self.colors = [color0, color1, color2, color3]
        self.tracker_bbox = []
        self.net_dir = params['rpn_net_cfg']
        self.gt = None
        self.iou_list = []
        self.method_list = ['SiamRPN', 'KCF']
        if track_method not in self.method_list:
            print('No such fucking tracker you choose.')
            exit
        else:
            self.method = track_method

    def tracker_update(self, frame, is_keyframe=False, key_bbox=None, gt=None):
        
        """
        We use a two-sides match algorithm.
        If a key bbox is matched with an exist tracker, that means the tracker is matched with the key bbox too.
        So both unmatched key bbox and tracker should be removed.

        @ param frame : the frame to update
        @ param is_keyframe : When we use YOLO, set True
        @ param key_bbox : if None, we don't use YOLO to check the tracking result
        """
        self.gt = gt
        tracker_ret = True
        if is_keyframe is True:
            # If there exit some key bboxes
            if len(key_bbox) != 0:
                # Using this to get the matched tracker
                tracker_matched_id = np.array([])
                # Compute the distance between each tracker bbox and each key bbox, and match them
                for k_bbox in key_bbox:

                    tr_bbox_id = 0
                    is_matched = False
                    
                    for tr_bbox in self.tracker_bbox:
                        if tr_bbox != None and k_bbox != None and self.compute_rects_overlap(tr_bbox, k_bbox) > self.overlap_threshold:
                            tracker_matched_id = np.append(tracker_matched_id, tr_bbox_id)
                            is_matched = True
                            break
                        tr_bbox_id += 1

                    # Create a tracker for the key bbox that not matched with a exist tracker
                    if is_matched is False and k_bbox != None:
                        self.tracker_state = np.append(self.tracker_state, 1)
                        if self.method == 'KCF':
                            self.tracker_list.append(cv2.TrackerKCF_create())
                        elif self.method == 'SiamRPN':
                            self.tracker_list.append(TrackerSiamRPN(net_path=self.net_dir))
                        self.tracker_list[len(self.tracker_list) - 1].init(frame, k_bbox)
                        self.tracker_bbox.append(None)
                        tracker_matched_id = np.append(tracker_matched_id, len(self.tracker_state) - 1)
                
                # Remove the tracker that not matched with a key bbox
                if len(tracker_matched_id) > 0:
                    for tr_bbox_id in range(0, len(self.tracker_state)):
                        if tr_bbox_id not in tracker_matched_id:
                            self.tracker_state[tr_bbox_id] = -1

            # If there is no key bbox, we need remove all the trackers
            else:
                self.tracker_state[:] = -1
        
        # The only one step of updating and removing
        for tracker_id in range(0, len(self.tracker_list)):
            if self.method == 'SiamRPN':
                ret, bbox, loss = self.tracker_list[tracker_id].update(frame, get_loss=False)
            elif self.method == 'KCF':
                ret, bbox = self.tracker_list[tracker_id].update(frame)

            # If a tracker failed to track, we should remove it
            if ret is False:
                self.tracker_state[tracker_id] = -1
                tracker_ret = False
            else:
                self.tracker_bbox[tracker_id] = bbox
                self.draw_bbox(frame, self.tracker_bbox[tracker_id])
        
        self.remove_tracker()
        
        if self.gt != None:
            self.get_iou_list()
        
        return tracker_ret
        
    def remove_tracker(self):

        """
        Remove the tracker whitch the state is marked with -1
        """

        origin_list_len = len(self.tracker_state)
        # Use a loop to remove the tracker
        for i in range(0, origin_list_len):
            # If delete an element, the indicy after the removed element will change
            # so we use a real-time length checking to avoid this bug
            if i >= len(self.tracker_state):
                break

            if self.tracker_state[i] == -1:
                del self.tracker_list[i]
                self.tracker_state = np.delete(self.tracker_state, i)
                del self.tracker_bbox[i]

    def draw_bbox(self, frame, bbox):

        """
        @ param frame : Draw the bbox in this frame
        @ param bboxes : the bbox
        """

        color_id = np.random.randint(0, 4)
        pt1 = (int(bbox[0]), int(bbox[1]))
        pt2 = (int(bbox[0]) + int(bbox[2]), int(bbox[1]) + int(bbox[3]))
        cv2.rectangle(frame, pt1, pt2, self.colors[color_id], 3)
    
    def compute_rects_overlap(self, tr_bbox, key_bbox):

        """
        Compute the overlap-ratio between the tracker bbox and key bbox

        @ param tr_bbox : the bounding box of a tracker, (top_x, top_y, width, height)
        @ param key_bbox : the bounding box from YOLO, (top_x, top_y, width, height)
        @ return : the overlap-ratio of two bbox
        """

        tr_bbox_x, tr_bbox_y, tr_bbox_w, tr_bbox_h = tr_bbox[0], tr_bbox[1], tr_bbox[2], tr_bbox[3]
        key_bbox_x, key_bbox_y, key_bbox_w, key_bbox_h = key_bbox[0], key_bbox[1], key_bbox[2], key_bbox[3]
        
        if tr_bbox_x + tr_bbox_w < key_bbox_x:
            return 0
        if tr_bbox_y + tr_bbox_h < key_bbox_y:
            return 0
        if tr_bbox_x > key_bbox_x + key_bbox_w:
            return 0
        if tr_bbox_y > key_bbox_y + key_bbox_h:
            return 0
        
        overlap_x = min(tr_bbox_x + tr_bbox_w, key_bbox_x + key_bbox_w) - max(tr_bbox_x, key_bbox_x)
        overlap_y = min(tr_bbox_y + tr_bbox_h, key_bbox_y + key_bbox_h) - max(tr_bbox_y, key_bbox_y)
        s = overlap_x * overlap_y
        ratio = s / (tr_bbox_w * tr_bbox_h + key_bbox_w * key_bbox_h - s)
        return ratio

    def get_iou_list(self):
        if len(self.tracker_list) != 0:
            self.iou_list.append(iou([self.tracker_bbox[0][0], self.tracker_bbox[0][1], self.tracker_bbox[0][2], self.tracker_bbox[0][3]], 
                                     [self.gt[0], self.gt[1], self.gt[2], self.gt[3]]))
        else:
            self.iou_list.append(0)